from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List, Optional

mcp = FastMCP("SQLParse MCP Server")

"""
面向软件：python-sqlparse

MCP tools 列表：
1. sqlformat_command - 格式化 SQL 文件
   用途：使用 sqlformat 命令格式化 SQL 文件，支持多种格式化选项
   参数：
     - input_file: 输入 SQL 文件路径
     - output_file: 输出文件路径（可选）
     - keywords: 关键字大小写（upper/lower，可选）
     - identifiers: 标识符大小写（upper/lower，可选）
     - language: 语言模式（python/sql，可选）
     - strip_comments: 是否去除注释（布尔值，可选）
     - reindent: 是否重新缩进（布尔值，可选）
     - indent_width: 缩进宽度（整数，可选）
     - indent_after_first: 是否在首行后缩进（布尔值，可选）
     - indent_columns: 是否缩进列（布尔值，可选）
     - reindent_aligned: 是否对齐缩进（布尔值，可选）
     - use_space_around_operators: 是否在操作符周围使用空格（布尔值，可选）
     - wrap_after: 换行宽度（整数，可选）
     - comma_first: 是否逗号优先（布尔值，可选）
     - encoding: 文件编码（字符串，可选）
   返回结构：
     {
       "success": bool,
       "command": str,
       "exit_code": int,
       "stdout": str,
       "stderr": str
     }
"""

@mcp.tool()
def sqlformat_command(
    input_file: str,
    output_file: Optional[str] = None,
    keywords: Optional[str] = None,
    identifiers: Optional[str] = None,
    language: Optional[str] = None,
    strip_comments: Optional[bool] = None,
    reindent: Optional[bool] = None,
    indent_width: Optional[int] = None,
    indent_after_first: Optional[bool] = None,
    indent_columns: Optional[bool] = None,
    reindent_aligned: Optional[bool] = None,
    use_space_around_operators: Optional[bool] = None,
    wrap_after: Optional[int] = None,
    comma_first: Optional[bool] = None,
    encoding: Optional[str] = None
):
    """
    使用 sqlformat 命令格式化 SQL 文件
    
    参数：
      input_file: 输入 SQL 文件路径
      output_file: 输出文件路径（可选）
      keywords: 关键字大小写，可选值为 upper 或 lower
      identifiers: 标识符大小写，可选值为 upper 或 lower
      language: 语言模式，可选值为 python 或 sql
      strip_comments: 是否去除注释
      reindent: 是否重新缩进
      indent_width: 缩进宽度
      indent_after_first: 是否在首行后缩进
      indent_columns: 是否缩进列
      reindent_aligned: 是否对齐缩进
      use_space_around_operators: 是否在操作符周围使用空格
      wrap_after: 换行宽度
      comma_first: 是否逗号优先
      encoding: 文件编码
    
    返回 JSON 结构：
      {
        "success": 命令是否成功执行（布尔值）,
        "command": 执行的完整命令（字符串）,
        "exit_code": 命令退出码（整数）,
        "stdout": 标准输出（字符串）,
        "stderr": 标准错误输出（字符串）
      }
    """
    
    cmd = ["sqlformat"]
    
    if output_file:
        cmd.extend(["-o", output_file])
    
    if keywords:
        cmd.extend(["--keywords", keywords])
    
    if identifiers:
        cmd.extend(["--identifiers", identifiers])
    
    if language:
        cmd.extend(["--language", language])
    
    if strip_comments:
        cmd.append("--strip-comments")
    
    if reindent:
        cmd.append("--reindent")
    
    if indent_width is not None:
        cmd.extend(["--indent_width", str(indent_width)])
    
    if indent_after_first:
        cmd.append("--indent_after_first")
    
    if indent_columns:
        cmd.append("--indent_columns")
    
    if reindent_aligned:
        cmd.append("--reindent_aligned")
    
    if use_space_around_operators:
        cmd.append("--use_space_around_operators")
    
    if wrap_after is not None:
        cmd.extend(["--wrap_after", str(wrap_after)])
    
    if comma_first:
        cmd.append("--comma_first")
    
    if encoding:
        cmd.extend(["--encoding", encoding])
    
    cmd.append(input_file)
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        
        return {
            "success": result.returncode == 0,
            "command": " ".join(cmd),
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": "Command timed out after 30 seconds"
        }
    except Exception as e:
        return {
            "success": False,
            "command": " ".join(cmd),
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()